import numpy as np
from globals import * 
from numpy import sin, cos, exp, sqrt, pi, mean
from scipy import weave
from scipy.weave import converters

class CavitySection(object):
    '''Object containing Cavity and drift information. Default is to build benchmark linac but will also accept file'''
    def __init__(self,file=None, E_in=150e6, freq=704.4e6, phi_sdeg = -15.0, length=0, V_rf = 20.7e6, cavNum = 150, driftL=1.75, cavPeriod=np.ones(150)*150,TMat=None):
        if file:
            self._loadInfo(file,E_in)
        else:
            self.freq = np.ones(cavNum)*freq
            self.phi_sdeg = np.ones(cavNum)*phi_sdeg
            self.phi_s = self.phi_sdeg*pi/180.0
            self.length = np.ones(cavNum)*length
            self.V_rf = np.ones(cavNum)*V_rf
            self.cavNum = cavNum
            self.drifts= np.ones(cavNum+1)*driftL
            self.createW_i(E_in)
        self.TMat = TMat
        self.cavPeriod = cavPeriod
        self.Corr = np.zeros(self.cavNum)
        self.V_eff = 1*self.V_rf
        self.phi_eff = 1*self.phi_s
        self.P = sqrt(self.W_i*(self.W_i + 2*proton_mass))
        self.PRatio = self.P[:-1]/self.P[1:]
        self.alignErr = np.zeros(self.cavNum)
        self.xlim = 28e-3

    def _loadInfo(self,file,E_in):
        '''Load file containing linac information drift length, voltage, phase, cavity length, R/Q, geometric beta'''
        d = np.loadtxt(file)
        self.cavNum = d.shape[0]
        self.drifts = np.zeros(self.cavNum+1)
        self.drifts[1:-1] = np.diff(d[:,0])
        self.V_rf = d[:,1]*1e6
        self.phi_sdeg = d[:,2]
        self.phi_s = d[:,2]*pi/180.0
        try:
            self.freq = d[:,4]
            self.length = d[:,3]
            self.RoverQ = d[:,5]
            self.beta_G = d[:,6]
        except IndexError:
            self.length=np.zeros(self.cavNum)
            self.freq = d[:,3]
        self.drifts[0]+=self.length[0]/2.0
        self.drifts[-1]+=self.length[-1]/2.0
        self.createW_i(E_in)
        self.P = np.sqrt(self.W_i*(self.W_i + 2*proton_mass))
        self.PRatio=self.P[:-1]/self.P[1:]

    def createW_i(self,E_in):
        '''Create energy profile based on voltage and phase'''
        self.W_i = np.zeros(self.cavNum+1)
        self.W_i[0] = E_in
        for cav in xrange(self.cavNum):
            self.W_i[cav+1] =self.W_i[cav] +  self.V_rf[cav]*cos(self.phi_s[cav]) 

    def VoltageCorrect(self):
        '''Perform interpolation on voltage as a function of energy to get gradient for tracking non sync particles'''
        V = self.V_rf
        W = self.W_i
        self.Corr = Interpolate(W,V,self.cavPeriod)

    def __add__(self,Cavobj):
        for key in self.__dict__.keys():
            if key in Cavobj.__dict__.keys():
                if key == 'cavNum':
                    self.cavNum += Cavobj.cavNum
                else:
                    exec('lst = [self.'+key+',Cavobj.'+key+']')
                    setattr(self,key,np.concatenate(lst))

    def VRatio(self,beta,beta_0,N):
        '''Calculate voltage ratio as a function of voltage at optimum beta from Steve's tech note'''
        if N%2 == 1:
            V = -4*beta**2/(N*pi*(beta_0**2 - beta**2))*cos(beta_0*N*pi/(beta*2))/sin(N*pi/2.0)
        elif N%2 == 0:
            V = 4*beta*beta_0/(N*pi*(beta_0**2 - beta**2))*sin(beta_0*N*pi/(beta*2))/cos(N*pi/2.0)
            pass
        self.V_rf*=V
        self.V_eff=1*self.V_rf

    def loadField(self,file,DesignBeta,DesignGradient):
        '''Load field map from acdtool in order to integrate fields, used to test benchmark voltage correction interpolation''' 
        d = np.loadtxt(file,skiprows=2)
        try:
            self.z = np.array([d[:,2]]*self.cavNum)
            self.Ez = np.array([d[:,5]]*self.cavNum)
        except IndexError:
            self.z = np.array([d[:,0]]*self.cavNum)*1e-3
            self.Ez = np.array([d[:,1]]*self.cavNum)
        self.lengthMap = np.ones(self.cavNum)*(self.z[0,-1] - self.z[0,0])
        for cav in xrange(self.cavNum):
            self.Ez[cav] *= DesignGradient[cav]/IntET(DesignBeta[cav],self.freq[cav],self.Ez[cav],self.z[cav])

    def IntegrationTrack(self):
        '''Create energy profile based on integrating fieldmap'''
        for cav in xrange(self.cavNum):
            beta = WToBeta(self.W_i[cav])
            self.V_rf[cav] = IntET(beta,self.freq[cav],self.Ez[cav],self.z[cav])#*self.lengthMap[cav]/self.length[cav]
            self.W_i[cav+1] = self.W_i[cav] + self.V_rf[cav]*cos(self.phi_s[cav])
        self.V_eff = 1*self.V_rf
            
    def AddRFError(self,grad_err,phase_err,distr='Uniform',seed=None):
        '''Apply errors to voltage as a percentage and phase in degrees'''
        if distr=='Gaussian':
            self.V_eff*=normal(1,grad_err,self.V_eff.size)   
            self.phi_eff*=normal(1,phase_err,self.phi_eff.size)   
        elif distr=='Uniform':
            self.V_eff=self.V_rf*uniform(1-grad_err,1+ grad_err,self.V_eff.size, seed)   
            self.phi_eff=uniform(self.phi_sdeg- phase_err, self.phi_sdeg+ phase_err,self.phi_eff.size,seed)*np.pi/180.0   

    def AddAlignmentError(self,sig,seed=None,cut=None,distr='Gaussian'):
        '''Apply alignment errors to cavities'''
        if distr=='Gaussian':
            self.alignErr=normal(0,sig,self.alignErr.size,seed,cut)   
        elif distr=='Uniform':
            self.alignErr=uniform(-sig,sig,self.alignErr.size,seed)

class Optics(object):
    def __init__(self,file):
        self.T6D = extractTMatrix(file)
        self.cavs = np.flatnonzero(self.T6D[:,5,4]!=0.0)
        self.T2D = self.T6D[:,0:2,0:2]

    def buildCondensedMat(self,lengthArr):
        self.T=np.zeros((self.cavs.size+1,2,2))
        T = self.T2D
        cavs = self.cavs
        start=0
        lst=[np.concatenate([T[0:cavs[0]],driftMat(lengthArr[0]/2.0)])]
        for i in xrange(cavs.size-1):
            d=T[cavs[i]+1:cavs[i+1]]
            lst.append(np.concatenate([driftMat(lengthArr[i]/2.0),d,driftMat(lengthArr[i+1]/2.0)]))
        lst.append(np.concatenate([driftMat(lengthArr[-1]/2.0),T[cavs[-1]+1:]]))
        for i in xrange(len(lst)):
            self.T[i]=MatrixDotArray(lst[i])

class HOM(object):
    '''Object containing Higher order mode information such as
    frequency, Qexternal and R/Q (Linac definition)'''
    def __init__(self,numCavs,freq=1e9,Qex=1e6,RoverQ=0,nmodes=1,type='Monopole'):
        self.freq = np.ones((numCavs,nmodes))*freq
        self.centerfreq = freq
        self.nmodes = nmodes
        self.Qex = np.ones((numCavs,nmodes))*Qex
        self.RoverQ = np.ones((numCavs,nmodes))*RoverQ
        self.HOMV = np.zeros((numCavs,nmodes),dtype=complex)
        self.type = type
        self.fsig = None
        self.corrRQ = np.zeros(self.RoverQ.shape)

    def RQCorrect(self,W_i,period):
        for i in xrange(self.nmodes):
            self.corrRQ[:,i] = Interpolate(W_i, self.RoverQ[:,i], period)

    def HOMInd(self,bunchQ):
        '''Pre calculate voltage induced by particle'''
        if self.type == 'Monopole':
            self.Vind = pi*self.freq*self.RoverQ*bunchQ
        elif self.type == 'Dipole':
            self.Vind = 2 * (pi*self.freq)**2 * self.RoverQ*bunchQ / speed_of_light

    def HOMDecay(self,t):
        '''Decay voltages are time defined from t'''
        self.HOMV*=exp(-pi*self.freq*t/self.Qex + 1j*2*pi*self.freq*t)        

    def AddFrequencyError(self, fsig, mode=0, cavStart=0, cavEnd=-1,seed=None, cut=None):
        '''Apply gaussian frequency spread to all cavities'''
        self.fsig = fsig
        self.freq[cavStart:cavEnd,mode] = normal(self.freq[cavStart:cavEnd,mode], fsig, self.freq[cavStart:cavEnd,mode].size, seed, cut)

    def AddFrequencyErrorSundelin(self, f0, mode=0, cavStart=0, cavEnd=-1,seed=None, cut=None):
        '''Apply gaussian frequency spread to all cavities using sigma from Sundelin's formula'''
        fsig = np.abs(f0 - self.freq[cavStart,mode])*1.09e-3
        self.freq[cavStart:cavEnd,mode] = normal(self.freq[cavStart:cavEnd,mode], fsig, self.freq[cavStart:cavEnd,mode].size, seed, cut)

    def AddFrequencyErrorSundelinQuad(self, f0, mode=0, cavStart=0, cavEnd=-1,seed=None, cut=None):
        '''Apply gaussian frequency spread to all cavities using sigma from Sundelin's formula'''
        sig = np.abs(f0 - self.freq[cavStart,mode])*1.09e-3
        dif = 0.0038*self.freq[cavStart,mode]
        fsig = np.sqrt(sig**2 + dif**2)
        self.freq[cavStart:cavEnd,mode] = normal(self.freq[cavStart:cavEnd,mode], fsig, self.freq[cavStart:cavEnd,mode].size, seed, cut)

    def loadHomFreq(self,file,mode,cavStart,cavNum):
        import pickle
        f = open(file)
        d = pickle.load(f)
        f.close()
        if type(mode)==int:
            self.freq[cavStart:cavStart+cavNum]=d['freq'][mode]
        elif type(mode)==np.ndarray:
            for i in xrange(np.array(mode).size):
                self.freq[cavStart:cavStart+cavNum,i]=d['freq'][mode[i]]

    def loadHomRQ(self,file,mode,cavStart,cavNum):
        import pickle
        f = open(file)
        d = pickle.load(f)
        f.close()
        if type(mode)==int:
            self.RoverQ[cavStart:cavStart+cavNum]=d['RQ'][cavStart:cavStart+cavNum,mode]
        elif type(mode)==np.ndarray:
            for i in xrange(np.array(mode).size):
                self.RoverQ[cavStart:cavStart+cavNum,i]=d['RQ'][cavStart:cavStart+cavNum,mode[i]]
    
    def loadHomFREQandRQ(self,file,mode,cavStart,cavNum):
        import pickle
        f = open(file)
        d = pickle.load(f)
        f.close()
        if type(mode)==int:
            self.freq[cavStart:cavStart+cavNum]=d['freq'][mode]
            self.RoverQ[cavStart:cavStart+cavNum]=d['RQ'][cavStart:cavStart+cavNum,mode]
        elif type(mode)==np.ndarray:
            for i in xrange(np.array(mode).size):
                self.freq[cavStart:cavStart+cavNum,i]=d['freq'][mode[i]]
                self.RoverQ[cavStart:cavStart+cavNum,i]=d['RQ'][cavStart:cavStart+cavNum,mode[i]]

    def pulseNum(self,Tr):
        Td = 2*self.Qex/(2*pi*self.freq)
        n = -Td/Tr*np.log(0.01) +1
        return int(np.max(np.ceil(n)))
 
class Pulse(object):
    '''Object containing pulse information such as number of bunches,
    bunch current and time between bunches'''
    def __init__(self,numBunches, numCavs, tb, Ib):
        self.cavNum = numCavs 
        self.numBunches = numBunches
        self.dE = np.zeros((numBunches,numCavs+2),dtype='float32')
        self.dt = np.zeros((numBunches,numCavs+2),dtype='float32') 
        self.x = self.dt
        self.xp = self.dE
        self.Lost = np.ones(numBunches)
        self.tb = tb
        self.Ib = Ib
        self.dq = np.ones(numBunches)
        self.Ibsig = None

    @property
    def emittance(self):
        e = pi*sqrt(mean(self.dE[self.Lost!=0,-1]**2)*mean(self.dt[self.Lost!=0,-1]**2) - mean(self.dE[self.Lost!=0,-1]*self.dt[self.Lost!=0,-1])**2)
        return e
    
    @property
    def rmsdE(self):
        rms = np.std(self.dE[self.Lost!=0,-1]) 
        return rms

    @property
    def rmsdt(self):
        rms = np.std(self.dt[self.Lost!=0,-1]) 
        return rms

    def AddEnergyError(self, Esig, seed=None, cut=None):
        '''Create gaussian distributed energy injection error'''
        self.dE[:,0] = normal(0,Esig,self.dE.shape[0], seed, cut)

    def AddTimeError(self, tsig, seed=None, cut=None):
        '''Create gaussian distributed time injection error'''
        self.dt[:,0] = normal(0,tsig,self.dt.shape[0], seed, cut)
    
    def AddPositionError(self, xsig, seed=None, cut=None):
        '''Create gaussian distributed position injection error'''
        self.x[:,0] = normal(0,xsig,self.x.shape[0], seed, cut)

    def AddMomentumError(self, xpsig, seed=None, cut=None):
        '''Create gaussian distributed momentum injection error'''
        self.xp[:,0] = normal(0,xpsig,self.xp.shape[0], seed, cut)

    def AddCurrentError(self, Ibsig, seed=None, cut=None):
        '''Create gaussian distributed charge jitter'''
        self.Ibsig = Ibsig
        self.dq = self.Ib/normal(self.Ib,Ibsig,self.dq.size, seed, cut)

    def loadInjectionNoise(self,file):
        import pickle
        f = open(file)
        d = pickle.load(f)
        f.close()
        self.dt[:,0]=d['dt']
        self.dE[:,0]=d['dE']
    
    @property
    def bunchQ(self):
        return self.Ib*self.tb

    def plotPhaseSpace(self,cav,xl=(-5,5),yl=(-20,20)):
        import matplotlib.pyplot as plt
        plt.plot(self.dt[:,cav]*360*704e6,self.dE[:,cav]*1e-6,'o')
        plt.xlabel('d$\mathrm{\phi}$ [deg at 704MHz]',fontsize=20)
        plt.ylabel('dE [MeV]',fontsize=20)
        plt.xlim(xl)
        plt.ylim(yl)
        plt.show()

class Linac(object):
    '''Master class to hold cavity, pulse and hom information and then perform tracking'''
    def __init__(self, CavObj, PulseObj, HOMObj=None):
        self.Cavs = CavObj
        self.Pulse = PulseObj
        self.HOM = HOMObj
        self.driftNum = self.Cavs.drifts.size
        self.dtdE = -self.Cavs.drifts/(speed_of_light*proton_mass*(self.Cavs.W_i/proton_mass*(self.Cavs.W_i/proton_mass + 2.0))**(1.5))
   
    def info(self):
        '''Print Linac info'''
        print '\n'
        print '-'*10+'Linac Info'+'-'*10 
        print 'Number of Cavities:', int(self.Cavs.cavNum)
        print 'Operating Frequency:', (', ').join(['%.2f' %(i*1e-6) for i in list(set(self.Cavs.freq))]), 'MHz'
        print 'Input Energy:', self.Cavs.W_i[0]*1e-6, 'MeV'
        print 'Ouput Energy:', self.Cavs.W_i[-1]*1e-6,'MeV'
        print '\n'
        print '-'*10+'HOM Info'+'-'*10
        if self.HOM:
            print 'HOM Frequency:', self.HOM.centerfreq*1e-9, 'GHz'
            if self.HOM.fsig:
                print 'HOM Frequency Spread:','%.1f' %(self.HOM.fsig*1e-6), 'MHz'
            print 'HOM Qexternal:', '%1.0e' %self.HOM.Qex[0]
            print 'HOM R/Q:', self.HOM.RoverQ[0]
        else:
            print 'No HOMs present'
        print '\n'
        print '-'*10+'Pulse Info'+'-'*10
        print 'Number of bunches:', int(self.Pulse.numBunches)
        print 'Nominal Beam Current:', self.Pulse.Ib*1e3, 'mA'
        if self.Pulse.Ibsig:
            print 'Charge jitter:', self.Pulse.Ibsig/self.Pulse.Ib*100, '%'
        print 'Bunches spacing: %1.2f' %(self.Pulse.tb*1e9), 'ns'
        print '\n'

    def TrackSlow(self):
        self.dV = np.zeros(self.Cavs.V_rf.shape)
        corr = self.Cavs.V_eff/self.Cavs.V_rf
        for b in xrange(self.Pulse.numBunches):
            for i in xrange(self.Pulse.cavNum):
                self.Pulse.dt[b,i+1] = self.dtdE[i]*self.Pulse.dE[b,i] + self.Pulse.dt[b,i]
                beta = WToBeta(self.Cavs.W_i[i]+self.Pulse.dE[b,i])
                V = IntET(beta,self.Cavs.freq[i],self.Cavs.Ez[i],self.Cavs.z[i])
                self.dV[i] = V - self.Cavs.V_rf[i]
                
                self.Pulse.dE[b,i+1] = self.Pulse.dE[b,i] + V*np.cos(self.Cavs.phi_eff[i] +2*pi*self.Cavs.freq[i]*self.Pulse.dt[b,i+1])  - self.Cavs.V_rf[i]*cos(self.Cavs.phi_s[i])

        self.Pulse.dt[b,self.Cavs.cavNum+1] = self.dtdE[self.Cavs.cavNum]*self.Pulse.dE[b,self.Cavs.cavNum] + self.Pulse.dt[b,self.Cavs.cavNum]
        self.Pulse.dE[b,self.Cavs.cavNum+1] = self.Pulse.dE[b,self.Cavs.cavNum] 
    
    def Track(self,BunchErrors=None):
        '''Tracks pulse through Linac calculating longitudinal error'''
        dE = self.Pulse.dE
        dt = self.Pulse.dt
        dtdE = self.dtdE
        V_eff = self.Cavs.V_eff
        phi_eff = self.Cavs.phi_eff
        omega = 2*pi*self.Cavs.freq
        V_rf = self.Cavs.V_rf
        phi_s = self.Cavs.phi_s
        phi_sdeg = self.Cavs.phi_sdeg
        phi_ss = self.Cavs.phi_sdeg/360.0/self.Cavs.freq
        cavNum = self.Cavs.cavNum
        numBunches = self.Pulse.numBunches
        corr = self.Cavs.Corr
        lost = self.Pulse.Lost
        if self.HOM:
            q = self.Pulse.dq*self.Pulse.bunchQ
            homv = self.HOM.HOMV
            homOmega = 2*pi*self.HOM.freq
            Qex = self.HOM.Qex
            tb = self.Pulse.tb
            RQ = self.HOM.RoverQ
            corrRQ = self.HOM.corrRQ
            nmodes = self.HOM.nmodes
            self.TrackHOMFast(dE,dt,dtdE,V_eff,phi_eff,omega,V_rf,phi_s,numBunches,cavNum,homv,homOmega,q,tb,Qex,RQ,corr,corrRQ,nmodes,phi_ss,lost)
        else:
            if BunchErrors:
                V_eff=V_rf*uniform(1-BunchErrors[0],1+ BunchErrors[0],((numBunches,cavNum)))   
                phi=phi_sdeg.repeat(numBunches).reshape((cavNum,numBunches)).T
                phi_eff = uniform(phi-BunchErrors[1],phi+BunchErrors[1],phi.shape)*pi/180.0
                self.TrackErrors(dE,dt,dtdE,V_eff,phi_eff,omega,V_rf,phi_s,numBunches,cavNum,corr,phi_ss,lost)
            else:
                self.TrackFast(dE,dt,dtdE,V_eff,phi_eff,omega,V_rf,phi_s,numBunches,cavNum,corr,phi_ss,lost)
        if int(self.Pulse.numBunches - np.flatnonzero(self.Pulse.Lost).size) != 0: 
            print 'Number of bunches lost:', int(self.Pulse.Lost.size - np.flatnonzero(self.Pulse.Lost).size)

    def TrackTrans(self):
        x = self.Pulse.x
        xp = self.Pulse.xp
        T = self.Cavs.TMat
        P = self.Cavs.P
        PRatio = self.Cavs.PRatio
        numBunches = self.Pulse.numBunches
        cavNum = self.Cavs.cavNum
        xlim = self.Cavs.xlim
        lost = self.Pulse.Lost
        if self.HOM:
            self.HOM.HOMInd(self.Pulse.bunchQ)
            homv = self.HOM.HOMV
            homOmega = 2*pi*self.HOM.freq
            vind = self.HOM.Vind 
            tb = self.Pulse.tb
            Qex = self.HOM.Qex
            dq = self.Pulse.dq
            nmodes = self.HOM.nmodes
            alignErr = self.Cavs.alignErr
            self.TrackTransHOMFast(x,xp,T,P,PRatio,numBunches,cavNum,homv,homOmega,vind, dq, tb, Qex,nmodes,alignErr,xlim,lost)
        else:
            self.TrackTransFast(x,xp,T,PRatio,numBunches,cavNum,xlim,lost)
        if int(self.Pulse.numBunches - np.flatnonzero(self.Pulse.Lost).size) != 0: 
            print 'Number of bunches lost:', int(self.Pulse.numBunches - np.flatnonzero(self.Pulse.Lost).size)

    def TrackErrors(self,dE,dt,dtdE,V_eff,phi_eff,omega,V_rf,phi_s,numBunches,cavNum,corr,phi_ss,lost):
        code = '''
        for (int b=0;b<numBunches;b++){
            for (int i=0;i<cavNum;i++){
                dt(b,i+1) = dtdE(i)*dE(b,i) + dt(b,i);
                if (dt(b,i+1)>(-2*phi_ss(i)) || dt(b,i+1)<(phi_ss(i))) {
                    lost(b)=0;
                    break;
                }
                dE(b,i+1) = dE(b,i) + (V_eff(b,i)+dE(b,i)*corr(i))*cos(phi_eff(b,i)+omega(i)*dt(b,i+1)) - V_rf(i)*cos(phi_s(i));
            }
            dt(b,cavNum+1) = dtdE(cavNum)*dE(b,cavNum) + dt(b,cavNum);
            dE(b,cavNum+1) = dE(b,cavNum);
        }
        '''
        weave.inline(code,['corr','numBunches','cavNum','dtdE','dE','dt','V_eff','phi_eff','omega','V_rf','phi_s','phi_ss','lost'],
                type_converters = converters.blitz)

    def TrackFast(self,dE,dt,dtdE,V_eff,phi_eff,omega,V_rf,phi_s,numBunches,cavNum,corr,phi_ss,lost):
        code = '''
        for (int b=0;b<numBunches;b++){
            for (int i=0;i<cavNum;i++){
                dt(b,i+1) = dtdE(i)*dE(b,i) + dt(b,i);
                if (dt(b,i+1)>(-2*phi_ss(i)) || dt(b,i+1)<(phi_ss(i))) {
                    lost(b)=0;
                    break;
                }
                dE(b,i+1) = dE(b,i) + (V_eff(i)+dE(b,i)*corr(i))*cos(phi_eff(i)+omega(i)*dt(b,i+1)) - V_rf(i)*cos(phi_s(i));
            }
            dt(b,cavNum+1) = dtdE(cavNum)*dE(b,cavNum) + dt(b,cavNum);
            dE(b,cavNum+1) = dE(b,cavNum);
        }
        '''
        weave.inline(code,['corr','numBunches','cavNum','dtdE','dE','dt','V_eff','phi_eff','omega','V_rf','phi_s','phi_ss','lost'],
                type_converters = converters.blitz)

    def TrackHOMFast(self,dE,dt,dtdE,V_eff,phi_eff,omega,V_rf,phi_s,numBunches,cavNum,homv,homOmega,q,tb,Qex,RQ,corr,corrRQ,nmodes,phi_ss,lost):
        j = 1j
        code = '''
        for (int b=0;b<numBunches;b++){
            for (int i=0;i<cavNum;i++){
                dt(b,i+1) = dtdE(i)*dE(b,i) + dt(b,i);
                if (dt(b,i+1)>(-2*phi_ss(i)) || dt(b,i+1)<(phi_ss(i))) {
                    lost(b)=0;
                    break;
                }
                dE(b,i+1) = dE(b,i) + (V_eff(i)+dE(b,i)*corr(i))*cos(phi_eff(i)+omega(i)*dt(b,i+1)) - V_rf(i)*cos(phi_s(i));
                for (int h=0;h<nmodes;h++){
                    float psi = homOmega(i,h)*dt(b,i+1);
                    double vind = q(b)*homOmega(i,h)/2.0*(RQ(i,h) + dE(b,i)*corrRQ(i,h));
                    dE(b,i+1) += (-homv(i,h).real()*cos(psi) + homv(i,h).imag()*sin(psi) - 0.5*vind);
                    homv(i,h)+=(vind*cos(psi) - j*vind*sin(psi));
                    homv(i,h)*=exp(-0.5*homOmega(i,h)*tb/Qex(i,h) + j*homOmega(i,h)*tb);
                }
            }
            dt(b,cavNum+1) = dtdE(cavNum)*dE(b,cavNum) + dt(b,cavNum);
            dE(b,cavNum+1) = dE(b,cavNum);
        }
        '''

        weave.inline(code,['j','numBunches','cavNum','dtdE','dE','dt','V_eff','phi_eff','omega','V_rf','phi_s',
                            'homOmega', 'homv', 'q','tb','Qex','RQ','corr','corrRQ','phi_ss','lost','nmodes'],
                type_converters = converters.blitz)

    def TrackTransFast(self,x,xp,T,PRatio,numBunches,cavNum,xlim,lost):
        code = '''
        for (int b=0;b<numBunches;b++){
            for (int i=0;i<cavNum;i++){
                x(b,i+1)= x(b,i)*T(i,0,0) + xp(b,i)*T(i,0,1);
                if (x(b,i+1)>xlim || x(b,i+1)<(-1*xlim)){
                    lost(b)=0;
                    break;
                }

                xp(b,i+1) = x(b,i)*T(i,1,0) + xp(b,i)*T(i,1,1);
                xp(b,i+1)*=PRatio(i);
            }
            x(b,cavNum+1)= x(b,cavNum)*T(cavNum,0,0) + xp(b,cavNum)*T(cavNum,0,1);
            xp(b,cavNum+1) = x(b,cavNum)*T(cavNum,1,0) + xp(b,cavNum)*T(cavNum,1,1);
        }
        '''

        weave.inline(code,['numBunches','cavNum','x','xp','T','PRatio','xlim','lost'],
                type_converters = converters.blitz)

    def TrackTransHOMFast(self,x,xp,T,P,PRatio,numBunches,cavNum,homv,homOmega, vind, dq, tb, Qex,nmodes,alignErr,xlim,lost):
        j=1j
        code = '''
        for (int b=0;b<numBunches;b++){
            for (int i=0;i<cavNum;i++){
                x(b,i+1) = x(b,i)*T(i,0,0) + xp(b,i)*T(i,0,1);
                if (x(b,i+1)>xlim || x(b,i+1)<(-1*xlim)){
                    lost(b)=0;
                    break;
                }
                xp(b,i+1) = x(b,i)*T(i,1,0) + xp(b,i)*T(i,1,1);
                for (int h=0;h<nmodes;h++){
                    xp(b,i+1) += (homv(i,h).imag())/P(i+1);
                    homv(i,h)-=((x(b,i+1)+alignErr(i))*vind(i)*dq(b));
                    homv(i,h)*=exp(-0.5*homOmega(i,h)*tb/Qex(i,h) + j*homOmega(i,h)*tb);
                }
                xp(b,i+1)*=PRatio(i);
            }
            x(b,cavNum+1) = x(b,cavNum)*T(cavNum,0,0) + xp(b,cavNum)*T(cavNum,0,1);
            xp(b,cavNum+1) = x(b,cavNum)*T(cavNum,1,0) + xp(b,cavNum)*T(cavNum,1,1);
        }
        '''
        weave.inline(code,['j','numBunches','cavNum','x','xp','T','P','PRatio','homv','homOmega','vind','dq','tb','Qex','nmodes','alignErr','xlim','lost'],
                type_converters = converters.blitz)

